package database

import (
	set2 "abdis/datastruct/set"
	"abdis/interface/database"
	"abdis/interface/datastruct"
	"abdis/interface/redis"
	"abdis/lib/utils"
	"abdis/redis/protocol"
	"strconv"
	"strings"
)

func checkSetAndReplyValid(dict datastruct.Set, reply protocol.ErrorReply) redis.Reply {
	if reply != nil {
		return reply
	}
	if dict == nil {
		return protocol.MakeNullBulkReply()
	}
	return nil
}

func (db *DB) getAsSet(key string) (datastruct.Set, protocol.ErrorReply) {
	entity, exists := db.GetEntity(key)
	if !exists {
		return nil, nil
	}
	set, ok := entity.Data.(datastruct.Set)
	if !ok {
		return nil, protocol.MakeWrongTypeErrReply()
	}
	return set, nil
}

func (db *DB) getOrInitSet(key string) (set datastruct.Set, isNew bool, errReply protocol.ErrorReply) {
	set, errReply = db.getAsSet(key)
	if errReply != nil {
		return nil, false, errReply
	}
	isNew = false
	if set == nil {
		set = set2.MakeSet()
		db.PutEntity(key, &database.DataEntity{Data: set})
		isNew = true
	}
	return set, isNew, nil
}

func undoSetChange(db *DB, args [][]byte) []database.CmdLine {
	return rollbackSetMembers(db, string(args[0]), utils.BytesToStrings(args[1:])...)
}

func execSAdd(db *DB, args [][]byte) (redis.Reply, bool) {
	key := string(args[0])
	vals := utils.BytesToStrings(args[1:])

	set, _, errReply := db.getOrInitSet(key)
	if errReply != nil {
		return errReply, false
	}

	count := 0
	for _, val := range vals {
		count += set.Add(val)
	}
	return protocol.MakeIntReply(int64(count)), true
}

func execSRem(db *DB, args [][]byte) (redis.Reply, bool) {
	key := string(args[0])
	members := utils.BytesToStrings(args[1:])

	set, errReply := db.getAsSet(key)
	if reply := checkSetAndReplyValid(set, errReply); reply != nil {
		return reply, false
	}

	count := 0
	for _, member := range members {
		count += set.Remove(member)
	}
	if set.Len() == 0 {
		db.Remove(key)
	}
	return protocol.MakeIntReply(int64(count)), true
}

func execSCard(db *DB, args [][]byte) (redis.Reply, bool) {
	key := string(args[0])
	set, errReply := db.getAsSet(key)
	if reply := checkSetAndReplyValid(set, errReply); reply != nil {
		return reply, false
	}
	return protocol.MakeIntReply(int64(set.Len())), true
}

func execSScan(db *DB, args [][]byte) (redis.Reply, bool) {
	count := 10
	pattern := "*"
	if len(args) > 2 {
		for i := 2; i < len(args); i++ {
			arg := strings.ToLower(string(args[i]))
			if arg == "count" {
				count0, err := strconv.Atoi(string(args[i+1]))
				if err != nil {
					return protocol.MakeSyntaxErrReply(), false
				}
				count = count0
				i++
			} else if arg == "match" {
				pattern = string(args[i+1])
				i++
			} else {
				return protocol.MakeSyntaxErrReply(), false
			}
		}
	}
	key := string(args[0])
	// get entity
	set, errReply := db.getAsSet(key)
	if errReply != nil {
		return errReply, false
	}
	if set == nil {
		return protocol.MakeEmptyMultiBulkReply(), false
	}
	cursor, err := strconv.Atoi(string(args[1]))
	if err != nil {
		return protocol.MakeErrReply("ERR invalid cursor"), false
	}

	keysReply, nextCursor := set.Scan(cursor, count, pattern)
	if nextCursor < 0 {
		return protocol.MakeErrReply("Invalid argument"), false
	}

	result := make([]redis.Reply, 2)
	result[0] = protocol.MakeBulkReply([]byte(strconv.FormatInt(int64(nextCursor), 10)))
	result[1] = protocol.MakeMultiBulkReply(keysReply)

	return protocol.MakeMultiRawReply(result), true
}

func init() {
	registerCommand("SAdd", execSAdd, writeFirstKey, undoSetChange, -3, flagWrite)
	registerCommand("SRem", execSRem, writeFirstKey, undoSetChange, -3, flagWrite)

	registerCommand("SCard", execSCard, readFirstKey, nil, 2, flagReadOnly)

	registerCommand("SScan", execSScan, readFirstKey, nil, -2, flagReadOnly)
}
